from numpy import poly1d, pi, sqrt
from sympy import Poly
from sympy.abc import x,y
from ThetaFunction import jacobiTheta
import sympy

class EllipticCurve:
    def __init__(self, a=None, b=None, isLattice=False):
        if a == None or b == None:
            self.a = sympy.abc.a
            self.b = sympy.abc.b
            return
        if isLattice == True:
            self.w1 = a
            self.w2 = b
            if ((self.w2/self.w1).imag < 0.0):
                self.w1 = b
                self.w2 = a

            self.tau = self.w2/self.w1
            
            g2 = (2.0/3.0) * (pi/self.w1)**4 * (jacobiTheta(1,0,0.0,self.tau)**8 + jacobiTheta(0,0,0.0,self.tau)**8 + jacobiTheta(0,1,0.0,self.tau)**8)
            g3 = (pi/self.w1)**6 * ((8.0/27.0) * (jacobiTheta(1,0,0.0,self.tau)**12 + jacobiTheta(0,0,0.0,self.tau)**12) - (4.0/9.0)*(jacobiTheta(1,0,0.0,self.tau)**4 + jacobiTheta(0,0,0.0,self.tau)**4) * jacobiTheta(1,0,0.0,self.tau)**4 * jacobiTheta(0,0,0.0,self.tau)**4)
            self.a = (-1.0/4.0) * g2
            self.b = (-1.0/4.0) * g3
        else:
            self.a = a
            self.b = b
            roots = sorted(poly1d([1,0,a,b]).r)
            self.e1 = roots[0]
            self.e2 = roots[1]
            self.e3 = roots[2]
             
            self.w1 = pi * (0.0+1.0j) / \
                self.arithGeomMean(sqrt(self.e3-self.e1),\
                                   sqrt(self.e2-self.e1))
            self.w2 = pi / \
                self.arithGeomMean(sqrt(self.e3-self.e1),\
                                   sqrt(self.e3-self.e2))
            if (self.w2/self.w1).imag < 0:
                temp = self.w2
                self.w2 = self.w1
                self.w1 = temp
            self.tau = self.w2/self.w1
                    
    def arithGeomMean(self,a,b):
        anew = a
        bnew = b
        tol = 0.0000000000001
        while(anew - bnew >= tol):
            aprev = anew
            bprev = bnew
            anew = 0.5 * (aprev+bprev)
            bnew = sqrt(aprev*bprev)
        return anew

    def divPoly(self, n):
        if n%2 == 0:
            return self.divPolyHelp(n)*Poly([1,0],y)
        else:
            return self.divPolyHelp(n)
        
    def divPolyHelp(self,n):
        if n == 1:
            return(Poly([1], x))
        elif n == 2:
            return(Poly([2], x))
        elif n == 3:
            return(Poly([3, 0, 6*self.a, 12*self.b, -self.a**2], x))
        elif n == 4:
            return(Poly([4,0, 20*self.a, 80*self.b, \
                               -20*self.a**2, -16*self.a*self.b, \
                               -32*self.b**2 - 4*self.a**3], x))
        elif n%2 == 0:
            m = n/2
            return(Poly(self.divPolyHelp(m) * (self.divPolyHelp(m+2)*\
                        self.divPolyHelp(m-1)**2 - self.divPolyHelp(m-2)*\
                        self.divPolyHelp(m+1)**2)/2))
        elif n%4 == 1:
            l = n/4
            return(Poly([1,0,self.a,self.b], x)**2 * \
                   self.divPolyHelp(2*l + 2) * self.divPolyHelp(2*l)**3 - \
                   self.divPolyHelp(2*l-1)*self.divPolyHelp(2*l+1)**3)
        elif n%4 == 3:
            l = n/4
            return(self.divPolyHelp(2*l+3)*self.divPolyHelp(2*l+1)**3 - \
                   Poly([1,0,self.a,self.b], x)**2 * \
                   self.divPolyHelp(2*l)*self.divPolyHelp(2*l+2)**3)
        
    def discrim(self):
        return -16*(4*self.a**3 + 27*self.b**2)
    
    def secondDeriv(self, f):
        return 4*(f.diff(x).diff(x) * \
                  Poly([1, 0, self.a, self.b], x) \
                  + f.diff(x) * Poly([3, 0, self.a], x)/2.0)
        
    def pFunc(self, z):
        z = z/self.w1
        val = pi**2 * jacobiTheta(0,0,0.0,self.tau)**2 * \
            jacobiTheta(1,0,0.0,self.tau)**2 * \
            jacobiTheta(0,1,z,self.tau)**2 / \
            jacobiTheta(1,1,z,self.tau)**2 - pi**2/3.0 * \
            (jacobiTheta(0,0,0.0,self.tau)**4 + \
             jacobiTheta(1,0,0.0,self.tau)**4)
        return val/(self.w1**2)
